import torch
import torch.nn as nn
import torch.nn.functional as F
from training.models.AttnSleep import MRCNN

class GradientReversalLayer(torch.autograd.Function):
    """梯度反转层，用于DANN中的领域对抗训练"""
    
    @staticmethod
    def forward(ctx, x, lambda_val=1.0):
        ctx.lambda_ = lambda_val
        return x.view_as(x)
    
    @staticmethod
    def backward(ctx, grad_output):
        # 添加数值稳定性措施
        grad_input = -ctx.lambda_ * grad_output
        # 添加梯度裁剪以防止梯度爆炸
        grad_input = torch.clamp(grad_input, -1.0, 1.0)
        return grad_input, None


class DomainAdaptiveAttnSleepWithGRL(nn.Module):
    """使用梯度反转层的领域自适应AttnSleep模型"""
    
    def __init__(self, num_classes=5, num_domains=2, **kwargs):
        super().__init__()
        # 使用标准MRCNN特征提取器
        self.feature_extractor = MRCNN(
            afr_reduced_cnn_size=kwargs.get('afr_reduced_cnn_size', 30)
        )
        
        # 获取特征维度
        features_dim = 30  # 这个值需要根据实际情况调整
        
        # 主任务分类器 - 睡眠阶段预测
        self.task_classifier = nn.Linear(features_dim, num_classes)
        
        # 领域分类器 - 用于对抗训练
        self.domain_classifier = nn.Sequential(
            nn.Linear(features_dim, 64),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(32, num_domains)
        )
        
        # 梯度反转层
        self.grl = GradientReversalLayer.apply
    
    def forward(self, x, lambda_=1.0):
        # 提取共享特征
        features = self.feature_extractor(x)
        
        # 如果特征是3D的 (batch_size, channels, seq_len)，需要进行池化
        if features.dim() == 3:
            features = torch.mean(features, dim=2)
        
        # 主任务预测
        task_pred = self.task_classifier(features)
        
        # 应用梯度反转层并进行领域分类
        reversed_features = self.grl(features, lambda_)
        domain_pred = self.domain_classifier(reversed_features)
        
        return task_pred, domain_pred